{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, numpy, json, os\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "from netdissect.progress import verbose_progress, default_progress\n",
    "from netdissect.nethook import edit_layers\n",
    "from netdissect.modelconfig import create_instrumented_model\n",
    "from netdissect.zdataset import standard_z_sample, z_sample_for_model\n",
    "from netdissect.easydict import EasyDict\n",
    "from netdissect.aceoptimize import ace_loss\n",
    "from netdissect.segmenter import UnifiedParsingSegmenter\n",
    "from netdissect.fullablate import measure_full_ablation\n",
    "from netdissect.plotutil import plot_tensor_images, plot_max_heatmap\n",
    "import netdissect.aceoptimize\n",
    "import netdissect.fullablate\n",
    "\n",
    "verbose_progress(True)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_ablation(scene, classname):\n",
    "    layer = 'layer4'\n",
    "    dissectdir = 'dissect/%s' % scene\n",
    "    with open(os.path.join(dissectdir, 'dissect.json')) as f:\n",
    "        dissection = EasyDict(json.load(f))\n",
    "\n",
    "    segmenter = UnifiedParsingSegmenter(segsizes=[256], segdiv='quad')\n",
    "    model = create_instrumented_model(dissection.settings)\n",
    "\n",
    "    for category in ['object', 'material', 'part']:\n",
    "        try:\n",
    "            classnum = segmenter.get_label_and_category_names()[0].index((classname, category))\n",
    "            break\n",
    "        except:\n",
    "            continue\n",
    "\n",
    "    edit_layers(model, [layer])\n",
    "\n",
    "    unit_count = 20\n",
    "\n",
    "    # then make a batch of 10 images\n",
    "    big_sample = z_sample_for_model(model, 1000, seed=3)\n",
    "    big_dataset = TensorDataset(big_sample)\n",
    "    big_loader = DataLoader(big_dataset, batch_size=10)\n",
    "\n",
    "\n",
    "\n",
    "    lrec = [l for l in dissection.layers if l.layer == layer][0]\n",
    "    rrec = [r for r in lrec.rankings if r.name == '%s-iou' % classname][0]\n",
    "    iou_scores = -torch.tensor(rrec.score)\n",
    "    iou_values, iou_order = (-iou_scores).sort(0)\n",
    "    iou_values = -iou_values\n",
    "    iou_ablation = torch.zeros_like(iou_scores)\n",
    "    iou_ablation[iou_order[:unit_count]] = 1\n",
    "\n",
    "    # load ablation from the tree model snapshot\n",
    "    snapdir = os.path.join(dissectdir, layer, 'ace', classname, 'snapshots')\n",
    "    data = torch.load(os.path.join(snapdir, 'epoch-9.pth'))\n",
    "    learned_scores = data['ablation'][0,:,0,0]\n",
    "    _, learned_order = (-learned_scores - iou_scores.cuda() * 1e-5).sort(0)\n",
    "    learned_values = learned_scores[learned_order]\n",
    "    learned_ablation = torch.zeros_like(learned_scores)\n",
    "    learned_ablation[learned_order[:unit_count]] = learned_scores[learned_order[:unit_count]]\n",
    "\n",
    "    progress = default_progress()\n",
    "\n",
    "    # (1) call ace_loss to get baseline tree pixels in batch\n",
    "    with torch.no_grad():\n",
    "        baseline_loss = 0\n",
    "        for [small_sample] in progress(big_loader):\n",
    "            baseline_loss += ace_loss(segmenter, classnum, model, layer,\n",
    "                    torch.zeros_like(learned_scores)[None,:,None,None],\n",
    "                    torch.zeros_like(learned_scores)[None,:,None,None],\n",
    "                    small_sample, 0, 0, 0, run_backward=False,\n",
    "                    discrete_pixels=True,\n",
    "                    discrete_units=0,\n",
    "                    # mixed_units=True,\n",
    "                    ablation_only=True,\n",
    "                    fullimage_measurement=True,\n",
    "                    fullimage_ablation=True)\n",
    "\n",
    "    # (2) apply 20 unit iou ablation and call ace_loss to see the difference\n",
    "    with torch.no_grad():\n",
    "        iou_loss = 0\n",
    "        for [small_sample] in progress(big_loader):\n",
    "            iou_loss += ace_loss(segmenter, classnum, model, layer,\n",
    "                    torch.zeros_like(iou_scores)[None,:,None,None].cuda(), # high_replacement\n",
    "                    iou_scores[None,:,None,None].cuda(),  # ablation\n",
    "                    small_sample, 0, 0, 0, run_backward=False,\n",
    "                    discrete_pixels=True,\n",
    "                    discrete_units=20,\n",
    "                    ablation_only=True,\n",
    "                    fullimage_measurement=True,\n",
    "                    fullimage_ablation=True)\n",
    "\n",
    "    # (3) apply 20 unit learned ablation and call ace_loss to see the difference\n",
    "    with torch.no_grad():\n",
    "        learned_loss = 0\n",
    "        for [small_sample] in progress(big_loader):\n",
    "            learned_loss += ace_loss(segmenter, classnum, model, layer,\n",
    "                    torch.zeros_like(learned_scores)[None,:,None,None].cuda(),\n",
    "                    learned_scores[None,:,None,None].cuda(),\n",
    "                    small_sample, 0, 0, 0, run_backward=False,\n",
    "                    discrete_pixels=True,\n",
    "                    discrete_units=20,\n",
    "                    mixed_units=True,\n",
    "                    ablation_only=True,\n",
    "                    fullimage_measurement=True,\n",
    "                    fullimage_ablation=True)\n",
    "\n",
    "    print('%s - %s' % (scene, classname))\n",
    "    print('baseline %g iou %g learned %g' % (baseline_loss.item(), iou_loss.item(), learned_loss.item()))\n",
    "    print('iou', 1 - iou_loss / baseline_loss)\n",
    "    print('learned', 1 - learned_loss / baseline_loss)\n",
    "    return (1 - learned_loss / baseline_loss).item()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# results = {}\n",
    "for key in [\n",
    "    ('conferenceroom', 'window'),\n",
    "    ('conferenceroom', 'chair'),\n",
    "    ('conferenceroom', 'curtain'),\n",
    "    ('conferenceroom', 'person'),\n",
    "    ('conferenceroom', 'table'),\n",
    "    ('churchoutdoor', 'window'),\n",
    "    ('diningroom', 'window'),\n",
    "    ('restaurant', 'window'),\n",
    "    ('kitchen', 'window'),\n",
    "    ('livingroom', 'window'),\n",
    "    ('bedroom', 'window'),\n",
    "    ('churchoutdoor', 'window'),\n",
    "    ('churchoutdoor', 'brick'),\n",
    "    ('churchoutdoor', 'cloud'),\n",
    "    ('churchoutdoor', 'dome'),\n",
    "    ('churchoutdoor', 'door'),\n",
    "    ('churchoutdoor', 'tree'),\n",
    "]:\n",
    "    if key not in results:\n",
    "        results[key] = evaluate_ablation(*key)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# {('conferenceroom', 'window'): 0.5683056116104126,\n",
    "#  ('conferenceroom', 'chair'): 0.12282001972198486,\n",
    "#  ('conferenceroom', 'curtain'): 0.5497019290924072,\n",
    "#  ('conferenceroom', 'person'): 0.614344596862793,\n",
    "#  ('conferenceroom', 'table'): 0.2871387004852295,\n",
    "#  ('churchoutdoor', 'window'): 0.4741958975791931,\n",
    "#  ('diningroom', 'window'): 0.3416675329208374,\n",
    "#  ('restaurant', 'window'): 0.46097666025161743,\n",
    "#  ('livingroom', 'window'): 0.26558423042297363,\n",
    "#  ('bedroom', 'window'): 0.2773325443267822,\n",
    "#  ('kitchen', 'window'): 0.5389571785926819,\n",
    "#  ('churchoutdoor', 'brick'): 0.427112877368927,\n",
    "#  ('churchoutdoor', 'cloud'): 0.27516114711761475,\n",
    "#  ('churchoutdoor', 'dome'): 0.7740037441253662,\n",
    "#  ('churchoutdoor', 'door'): 0.5343429446220398,\n",
    "#  ('churchoutdoor', 'tree'): 0.5980440378189087}\n",
    "\n",
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_ablation('churchoutdoor', 'tree')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}